from multiprocessing import Pool

import cv2
import numpy as np
import os
import pydicom as dicom
import pandas as pd
import shutil
import SimpleITK as sitk
from tqdm import tqdm
from sklearn.model_selection import train_test_split
from tool import getTime
def pixel_array_normalize(pixel_array):
    return 255.0 * (pixel_array - pixel_array.min()) / (pixel_array.max() - pixel_array.min())

def get_center_slice(annotation_path):
    """
    获取每个病例病灶中心切片号
    :param annotation_path: 标注文件路径
    :return: 字典 key为病例ID value为切片号
    """
    annotations = pd.read_excel(annotation_path)
    caseId2sliceId = {}
    for annotation in annotations.itertuples():
        case_id = annotation[1]
        start_slice = annotation[6]
        end_slice = annotation[7]
        caseId2sliceId[case_id] = int((end_slice - start_slice) / 2 + start_slice)
    return caseId2sliceId

def get_subtraction(array1, array2):
    """
    获取剪影图像
    :param dicom1:  增强后的dicom文件路径
    :param dicom2:  增强前的dicom文件路径
    :return:        剪影图像矩阵
    """
    pa1 = array1
    pa2 = array2
    pa1 = pixel_array_normalize(pa1)
    pa2 = pixel_array_normalize(pa2)
    img = pa1 - pa2

    # 截断负值
    img = np.clip(img, a_min=0, a_max=img.max())
    img = pixel_array_normalize(img)
    # 计算直方图
    img = np.uint8(img)
    hist = cv2.calcHist([img], [0], None, [255], [1, 256])
    # plt.plot(hist)
    # plt.show()
    hist_sum = np.cumsum(hist)
    # 截断归一化
    min_val = 0
    # min_val = np.searchsorted(hist_sum, hist_sum[-1] * 0.01, side='left')
    max_val = np.searchsorted(hist_sum, hist_sum[-1] * 0.9998, side='left')
    img = np.clip(img, a_min=min_val, a_max=max_val)
    img = 255.0 * (img - min_val) / (max_val - min_val)
    return img

# 获取duke数据集的剪影jpg图像

def duke2subtraction2d(file):
    nii2d_save_path = "/data1/home/liukai/AllData/Duke/Duke/nii/subtractSecond2D"
    nii_path = "/data1/home/liukai/AllData/Duke/Duke/nii"
    pre_dir = os.path.join(nii_path, "pre")
    post2_dir = os.path.join(nii_path, "post2")

    pre_nii = sitk.ReadImage(os.path.join(pre_dir, file))
    pre_array = sitk.GetArrayFromImage(pre_nii).astype(np.int32)

    post2_nii = sitk.ReadImage(os.path.join(post2_dir, file))
    post2_array = sitk.GetArrayFromImage(post2_nii).astype(np.int32)

    #-------训练集裁剪中心三张图片-----#
    file_name_list=[]
    select_ids = [i for i in range(0,post2_array.shape[0])]
    for id in select_ids:
        # 因为数组下标是从0开始，sliceId下标是从1开始，所以对应的数组位置要减1
        pre_center_array = pre_array[id, ::]
        post2_center_array = post2_array[id, ::]
        img_array = get_subtraction(pre_center_array, post2_center_array)
        file_name=file.split('.')[0] + '_second_subtraction_norm_' + str(id) + '.nii.gz'
        file_name_list.append(file_name)

        nii_file = sitk.GetImageFromArray(img_array)
        sitk.WriteImage(nii_file, os.path.join(nii2d_save_path, file_name))
    print("该文件有{}张切片".format(len(file_name_list)))
    return len(file_name_list)








if __name__ == '__main__':

    #-----------------------把3D的nii图像进行切片，切成2D的nii.gz-----------------------#
    nii_path="/data1/home/liukai/AllData/Duke/Duke/nii"
    import time

    start = time.time()
    file_name_list = os.listdir(os.path.join(nii_path, 'pre'))
    print("总共有{}个3nii文件".format(len(file_name_list)))
    total_file_number=0
    for file in tqdm(file_name_list):
        l=duke2subtraction2d(file)
        total_file_number+=l
    end=time.time()
    print("切片共{}张,花费时间:{}/min".format(str(total_file_number),(end-start)/60))
